#include <iostream>
#include <vector>
#include <random>
#include <stdexcept>
#include <algorithm>
#include <openssl/sha.h>
#include <gmp.h>
#include <gmpxx.h>

using namespace std;

// 全局参数: 素数 p = 2^128 + 51
const mpz_class P = (mpz_class(1) << 128) + 51;

// 辅助函数：将mpz_class转换为字节向量
vector<uint8_t> mpz_to_bytes(const mpz_class& num, size_t len = 16) {
    vector<uint8_t> bytes(len, 0);
    mpz_class n = num;
    
    // 导出大整数到字节数组
    size_t count;
    mpz_export(bytes.data(), &count, 1, 1, 0, 0, n.get_mpz_t());
    
    // 如果需要，反转字节顺序为大端序
    if (count < len) {
        reverse(bytes.begin(), bytes.begin() + count);
    }
    
    return bytes;
}

// 辅助函数：将字节向量转换为mpz_class
mpz_class bytes_to_mpz(const vector<uint8_t>& bytes) {
    mpz_class result;
    // 导入字节数组到大整数
    mpz_import(result.get_mpz_t(), bytes.size(), 1, 1, 0, 0, bytes.data());
    return result;
}

// 模逆元计算
mpz_class mod_inv(const mpz_class& a, const mpz_class& m) {
    mpz_class inv;
    if (mpz_invert(inv.get_mpz_t(), a.get_mpz_t(), m.get_mpz_t()) == 0) {
        throw runtime_error("mod_inv does not exist!");
    }
    return inv;
}

// 密钥派生函数 (基于SHA-256)
vector<uint8_t> kdf(const mpz_class& key, size_t len) {
    vector<uint8_t> stream;
    // 将主密钥转换为固定长度字节（16字节）
    vector<uint8_t> key_bytes = mpz_to_bytes(key, 16);
    uint32_t ctr = 0; // 计数器

    while (stream.size() < len) {
        // 构造输入：密钥字节 + 计数器（4字节，大端序）
        vector<uint8_t> input = key_bytes;
        for (int i = 3; i >= 0; --i) {
            input.push_back(static_cast<uint8_t>((ctr >> (i*8)) & 0xFF));
        }

        // 计算SHA-256
        uint8_t hash[SHA256_DIGEST_LENGTH];
        SHA256(input.data(), input.size(), hash);
        
        // 将哈希值添加到流
        size_t needed = len - stream.size();
        size_t n = min(static_cast<size_t>(SHA256_DIGEST_LENGTH), needed);
        stream.insert(stream.end(), hash, hash + n);

        // 更新计数器
        ctr++;
    }
    return stream;
}

// 加密函数
pair<vector<uint8_t>, vector<pair<int, mpz_class>>> 
encrypt(const string& plaintext, int share_num, int limit_num) {
    // 初始化GMP随机状态
    gmp_randstate_t state;
    gmp_randinit_default(state);
    
    // 使用系统随机设备作为种子
    random_device rd;
    gmp_randseed_ui(state, rd());

    // 生成128位主密钥 (0 到 2^128-1)
    mpz_class K;
    mpz_urandomb(K.get_mpz_t(), state, 128);
    
    // 随机多项式系数
    mpz_class a[limit_num-1];
    for(int i = 0; i < limit_num - 1; i++)
        mpz_urandomm(a[i].get_mpz_t(), state, P.get_mpz_t());
    
    // 清理随机状态
    gmp_randclear(state);

    // 计算密钥份额 (x=1,2,3,4)
    vector<pair<int, mpz_class>> shares;
    for (int x = 1; x <= share_num; ++x) {
        mpz_class fx = K % P;
        for(int i = 0; i< limit_num - 1; i++){
            mpz_class pow_tmp;
            mpz_class base = x;
            // mpz_init(base);
            // mpz_init(pow_tmp);
            // mpz_set_ui(base, x);
            mpz_pow_ui(pow_tmp.get_mpz_t(), base.get_mpz_t(), i+1);
            fx = (fx + a[i]*pow_tmp) % P;
        }
        if (fx < 0) fx += P;
        shares.push_back({x, fx});
    }
    
    // 加密明文
    vector<uint8_t> key_stream = kdf(K, plaintext.size());
    vector<uint8_t> ciphertext;
    for (size_t i = 0; i < plaintext.size(); ++i) {
        ciphertext.push_back(plaintext[i] ^ key_stream[i]);
    }
    
    return {ciphertext, shares};
}

// 解密函数
string decrypt(const vector<uint8_t>& ciphertext, 
               const vector<pair<int, mpz_class>>& shares) {
    if (shares.size() < 3) {
        throw runtime_error("At least three secret shares are needed!");
    }
    
    // 拉格朗日插值恢复K
    const auto& [x1, y1] = shares[0];
    const auto& [x2, y2] = shares[1];
    const auto& [x3, y3] = shares[2];
    
    // 计算分母
    mpz_class denom1 = (x1 - x2) * (x1 - x3);
    mpz_class denom2 = (x2 - x1) * (x2 - x3);
    mpz_class denom3 = (x3 - x1) * (x3 - x2);
    
    // 计算基函数在0处的值
    mpz_class L1 = (x2 * x3) * mod_inv(denom1, P) % P;
    mpz_class L2 = (x1 * x3) * mod_inv(denom2, P) % P;
    mpz_class L3 = (x1 * x2) * mod_inv(denom3, P) % P;
    
    mpz_class K = (y1*L1 + y2*L2 + y3*L3) % P;
    if (K < 0) K += P;
    
    // 解密密文
    vector<uint8_t> key_stream = kdf(K, ciphertext.size());
    string plaintext;
    for (size_t i = 0; i < ciphertext.size(); ++i) {
        plaintext.push_back(ciphertext[i] ^ key_stream[i]);
    }
    return plaintext;
}

// 打印十六进制数据
void print_hex(const vector<uint8_t>& data) {
    for (uint8_t c : data) {
        printf("%02X", c);
    }
    cout << endl;
}

// 打印密钥份额
void print_shares(const vector<pair<int, mpz_class>>& shares) {
    for (const auto& [x, y] : shares) {
        cout << "Share " << x << ": ";
        vector<uint8_t> bytes = mpz_to_bytes(y);
        print_hex(bytes);
    }
}

void op_encrypt(){
    int limit_num, share_num;
    cout << "Limit: ";
    cin >> limit_num;
    cout << "Share: ";
    cin >> share_num;
    getchar();
    string plaintext;
    cout<<"Plaintext to be encrypted:"<<endl;
    getline(cin, plaintext);

    auto [ciphertext, shares] = encrypt(plaintext, share_num, limit_num);
    cout << "\nCiphertext (Hex): ";
    print_hex(ciphertext);
    
    cout << "\nShares:" << endl;
    print_shares(shares);

    // vector<pair<int, mpz_class>> sub_shares1(shares.begin(), shares.begin() + 3);
    // string decrypted1 = decrypt(ciphertext, sub_shares1);
    // cout << "\nUsing shares 1-3 for decryption: " << decrypted1 << endl;

    // vector<pair<int, mpz_class>> sub_shares2(shares.begin()+2, shares.begin() + 5);
    // string decrypted2 = decrypt(ciphertext, sub_shares2);
    // cout << "\nUsing shares 3-5 for decryption: " << decrypted2 << endl;
    // print_shares(sub_shares2);

    return;
}

// void op_decrypt(){

//     cout<<"At least 3 secret shares are needed here."<<endl;
//     vector<uint8_t> ciphertext;
//     for (size_t i = 0; i < plaintext.size(); ++i) {
//         ciphertext.push_back(getchar());
//     }

//     vector<pair<int, mpz_class>> shares;
//     int x, fx;
//     for (int i = 1; i <= 4; i++) {
//         cout<<"Secret share 1: "
//         cin >> x >> fx;
//         shares.push_back({x, fx});
//     }


// }

int main() {
    // string plaintext = "Hello, World! This is a (3,4) shamir demo, with GMP enabled.";
    // cout << "Plaintext: " << plaintext << endl;
    // cout << "Pltext length: " << plaintext.length() << " bytes" << endl;
    
    // // 加密
    // auto [ciphertext, shares] = encrypt(plaintext);
    // cout << "\nCiphertext (Hex): ";
    // print_hex(ciphertext);
    
    // cout << "\nShares:" << endl;
    // print_shares(shares);
    
    op_encrypt();
    // 测试1：使用前3个份额解密
    // vector<pair<int, mpz_class>> sub_shares1(shares.begin(), shares.begin() + 3);
    // string decrypted1 = decrypt(ciphertext, sub_shares1);
    // cout << "\nUsing shares 1-3 for decryption: " << decrypted1 << endl;

    
    return 0;
}
